/*
 * Copyright 2019 The TensorFlow Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.linkai.jacklinmap.ObjectDetection;

import android.annotation.SuppressLint;
import android.graphics.Bitmap;
import android.graphics.Bitmap.Config;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Matrix;
import android.graphics.Paint;
import android.graphics.Paint.Style;
import android.graphics.RectF;
import android.graphics.Typeface;
import android.media.ImageReader.OnImageAvailableListener;
import android.os.Build;
import android.os.Bundle;
import android.os.SystemClock;
import android.util.Log;
import android.util.Size;
import android.util.TypedValue;
import android.view.View;
import android.widget.RadioButton;
import android.widget.RadioGroup;
import android.widget.Toast;

import com.linkai.jacklinmap.ObjectDetection.Adapter.ObjectResultRecyclerViewAdapter;
import com.linkai.jacklinmap.ObjectDetection.customview.OverlayView;
import com.linkai.jacklinmap.ObjectDetection.customview.OverlayView.DrawCallback;
import com.linkai.jacklinmap.ObjectDetection.env.BorderedText;
import com.linkai.jacklinmap.ObjectDetection.env.ImageUtils;
import com.linkai.jacklinmap.ObjectDetection.env.Logger;
import com.linkai.jacklinmap.ObjectDetection.tflite.Classifier;
import com.linkai.jacklinmap.ObjectDetection.tflite.DetectorFactory;
import com.linkai.jacklinmap.ObjectDetection.tflite.YoloV5Classifier;
import com.linkai.jacklinmap.ObjectDetection.tracking.MultiBoxTracker;
import com.linkai.jacklinmap.R;

import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

import androidx.annotation.RequiresApi;

/**
 * 创建日期：2021/6/23 9:25
 * @author 林凯
 * 文件名称： DetectorActivity.java
 * 类说明： 继承 CameraActivity
 *        主要关注以下方法：
 *     （1）onCreate() 方法里面的 radioGroup 监听事件
 *              在这里面更换模型，之后对 YoloV5Classifier 对象进行更新。同时还要更新 RecycleView。
 *              涉及到很多线程的知识
 *     （2）onPreviewSizeChosen()
 *          进行模型的选择，用于检测的类 YoloV5Classifier 的实例化， 和一些 BitMap 的初始化。
 *     （3）processImage()
 *          主要进行图像的处理，调用 YoloV5Classifier 进行检测。
 *          检测完成之后使用 Canvas 绘制矩形框
 */
@RequiresApi(api = Build.VERSION_CODES.KITKAT)
public class DetectorActivity extends com.linkai.jacklinmap.ObjectDetection.CameraActivity implements OnImageAvailableListener {
    private static final Logger LOGGER = new Logger();

    private static final DetectorMode MODE = DetectorMode.TF_OD_API;

    // API 的最小置信度，只有大于这个置信度的目标才会显示出来
    private static final float MINIMUM_CONFIDENCE_TF_OD_API = 0.3f;

    private static final boolean MAINTAIN_ASPECT = true;

    // 想要预览的图像大小，就是展示图像和矩形框的界面的大小。
//    @SuppressLint("NewApi")
//    private static final Size DESIRED_PREVIEW_SIZE = new Size(640, 640);
    @SuppressLint("NewApi")
    private static final Size DESIRED_PREVIEW_SIZE = new Size(360, 640);


    private static final boolean SAVE_PREVIEW_BITMAP = false;
    private static final float TEXT_SIZE_DIP = 10;
    OverlayView trackingOverlay;
    private Integer sensorOrientation;

    private YoloV5Classifier detector;

    private long lastProcessingTimeMs;

    // 矩形框对应的 BitMap
    private Bitmap rgbFrameBitmap = null;

    // 手机摄像头捕捉的 BitMap
    private Bitmap croppedBitmap = null;

    //
    private Bitmap cropCopyBitmap = null;

    private boolean computingDetection = false;

    private long timestamp = 0;

    private Matrix frameToCropTransform;
    private Matrix cropToFrameTransform;

    private MultiBoxTracker tracker;

    private BorderedText borderedText;

    // 单选按钮组
    private RadioGroup radioGroup;

    // 单选按钮当前的选中情况; 初始情况默认选中通用物体识别
    String radioButtonSelect = "yolov5s-fp16.tflite";

    @RequiresApi(api = Build.VERSION_CODES.LOLLIPOP)
    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);

        // 单选按钮初始化
        radioGroup = findViewById(R.id.object_result_radiogroup);

        // 单选按钮添加点击事件。同时还要更新 YoloV5Classifier 对象
        radioGroup.setOnCheckedChangeListener(new RadioGroup.OnCheckedChangeListener() {
            @Override
            public void onCheckedChanged(RadioGroup group, int checkedId) {
                RadioButton radioButton = (RadioButton) findViewById(checkedId);
                String checked = radioButton.getText().toString();
                radioButtonSelect = checked;

                /*
                    checked 的值有2中情况
                    （1） yolov5s-fp16.tflite   通用物体检测； 对应的 radiobutton_01
                    （2）yolov5s_ae86_640-fp16.tflite   特定的4种物体检测；  对应的 radiobutton_02
                * */

                handler.post(new Runnable() {
                    @Override
                    public void run() {

                        // 再创建一个子线程，在子线程中操作 UI 界面
                        runOnUiThread(
                                new Runnable() {
                                    @Override
                                    public void run() {
                                        // 为 RecycleView 重新设置适配器，主要就是更换 Model
                                        objectResultRecyclerViewAdapter = new ObjectResultRecyclerViewAdapter(getApplicationContext());
                                        objectResultRecyclerViewAdapter.setModelType("yolov5s_ae86_640-fp16.tflite");
                                        ArrayList<Classifier.Recognition> results = new ArrayList<Classifier.Recognition>();
                                        objectResultRecyclerViewAdapter.setResult(results);
                                        recyclerView.setAdapter(objectResultRecyclerViewAdapter);
                                    }
                                });

                        try {
                            detector = DetectorFactory.getDetector(getAssets(), radioButtonSelect);
                            // Customize the interpreter to the type of device we want to use.
                            if (detector == null) {
                                return;
                            }
                        } catch (IOException e) {
                            e.printStackTrace();
                            LOGGER.e(e, "Exception in updateActiveModel()");
                            Toast toast =
                                    Toast.makeText(
                                            getApplicationContext(), "Classifier could not be initialized", Toast.LENGTH_SHORT);
                            toast.show();
                            finish();
                        }

                        int cropSize = detector.getInputSize();
                        croppedBitmap = Bitmap.createBitmap(cropSize, cropSize, Config.ARGB_8888);

                        frameToCropTransform =
                                ImageUtils.getTransformationMatrix(
                                        previewWidth, previewHeight,
                                        cropSize, cropSize,
                                        sensorOrientation, MAINTAIN_ASPECT);

                        cropToFrameTransform = new Matrix();
                        frameToCropTransform.invert(cropToFrameTransform);
                    }
                });

            }
        });
    }

    /**
     * 在 CameraActivity 里面的 onPreviewFrame回调方法里面调用
     */
    @RequiresApi(api = Build.VERSION_CODES.LOLLIPOP)
    @Override
    public void onPreviewSizeChosen(final Size size, final int rotation) {
        final float textSizePx =
                TypedValue.applyDimension(
                        TypedValue.COMPLEX_UNIT_DIP, TEXT_SIZE_DIP, getResources().getDisplayMetrics());
        borderedText = new BorderedText(textSizePx);
        borderedText.setTypeface(Typeface.MONOSPACE);

        tracker = new MultiBoxTracker(this);

        final String modelString = radioButtonSelect;

        // yolov5s-fp16.tflite
        /*
            yolov5s-fp16.tflite     识别 coco数据集的模型名称
            yolov5s_ae86-fp16.tflite        识别 “AE86”，“手指滑板”，“折扇”，“陀螺” 4种物体的模型名称
        * */
        System.out.println("modelString");
        System.out.println(modelString);

        try {
            detector = DetectorFactory.getDetector(getAssets(), modelString);
        } catch (final IOException e) {
            e.printStackTrace();
            LOGGER.e(e, "Exception initializing classifier!");
            Toast toast =
                    Toast.makeText(
                            getApplicationContext(), "Classifier could not be initialized", Toast.LENGTH_SHORT);
            toast.show();
            finish();
        }

        int cropSize = detector.getInputSize();

        previewWidth = size.getWidth();
        previewHeight = size.getHeight();

        sensorOrientation = rotation - getScreenOrientation();
        LOGGER.i("Camera orientation relative to screen canvas: %d", sensorOrientation);

        LOGGER.i("Initializing at size %dx%d", previewWidth, previewHeight);
        rgbFrameBitmap = Bitmap.createBitmap(previewWidth, previewHeight, Config.ARGB_8888);
        croppedBitmap = Bitmap.createBitmap(cropSize, cropSize, Config.ARGB_8888);

        frameToCropTransform =
                ImageUtils.getTransformationMatrix(
                        previewWidth, previewHeight,
                        cropSize, cropSize,
                        sensorOrientation, MAINTAIN_ASPECT);

        cropToFrameTransform = new Matrix();
        frameToCropTransform.invert(cropToFrameTransform);

        trackingOverlay = (OverlayView) findViewById(R.id.tracking_overlay);
        trackingOverlay.addCallback(
                new DrawCallback() {
                    @Override
                    public void drawCallback(final Canvas canvas) {
                        tracker.draw(canvas);
                        if (isDebug()) {
                            tracker.drawDebug(canvas);
                        }
                    }
                });

        tracker.setFrameConfiguration(previewWidth, previewHeight, sensorOrientation);
    }


    /**
     * 处理图像的关键代码（点击展开详细信息）；在 CameraActivity 里面的 onPreviewFrame回调方法里面调用
     * (1)
     */
    @Override
    protected void processImage() {
        ++timestamp;
        final long currTimestamp = timestamp;
        trackingOverlay.postInvalidate();

        // No mutex needed as this method is not reentrant.
        if (computingDetection) {
            readyForNextImage();
            return;
        }
        computingDetection = true;
        LOGGER.i("Preparing image " + currTimestamp + " for detection in bg thread.");

        rgbFrameBitmap.setPixels(getRgbBytes(), 0, previewWidth, 0, 0, previewWidth, previewHeight);

        readyForNextImage();

        final Canvas canvas = new Canvas(croppedBitmap);
        canvas.drawBitmap(rgbFrameBitmap, frameToCropTransform, null);
        // For examining the actual TF input.
        if (SAVE_PREVIEW_BITMAP) {
            ImageUtils.saveBitmap(croppedBitmap);
        }

        // 新开启一个线程来处理图像
        runInBackground(
                new Runnable() {
                    @Override
                    public void run() {
                        LOGGER.i("Running detection on image " + currTimestamp);
                        final long startTime = SystemClock.uptimeMillis();

                        // 获得图像的检测结果； 这个 croppedBitmap 就是手机摄像头捕捉到的画面
                        final List<Classifier.Recognition> results = detector.recognizeImage(croppedBitmap);


                        // 获取检测的时间
                        lastProcessingTimeMs = SystemClock.uptimeMillis() - startTime;

                        Log.e("CHECK", "run: " + results.size());

                        cropCopyBitmap = Bitmap.createBitmap(croppedBitmap);
                        final Canvas canvas = new Canvas(cropCopyBitmap);
                        final Paint paint = new Paint();
                        paint.setColor(Color.RED);
                        paint.setStyle(Style.STROKE);
                        paint.setStrokeWidth(2.0f);

                        // 设置最小置信度，只有大于这个置信度的目标才会被检测出来
                        float minimumConfidence = MINIMUM_CONFIDENCE_TF_OD_API;
                        switch (MODE) {
                            case TF_OD_API:
                                minimumConfidence = MINIMUM_CONFIDENCE_TF_OD_API;
                                break;
                        }

                        final List<Classifier.Recognition> mappedRecognitions =
                                new LinkedList<Classifier.Recognition>();

                        for (final Classifier.Recognition result : results) {
                            final RectF location = result.getLocation();

                            System.out.println("getId" + result.getId());
                            System.out.println("getTitle" + result.getTitle());
                            // 获得检测的类别
                            System.out.println("getConfidence" + result.getConfidence());
                            System.out.println("getDetectedClass" + result.getDetectedClass());


                            // 置信度大于最小置信度才显示
                            if (location != null && result.getConfidence() >= minimumConfidence) {

                                // 根据得到的结果，利用 canvas 绘制矩形框
                                canvas.drawRect(location, paint);

                                cropToFrameTransform.mapRect(location);

                                result.setLocation(location);
                                mappedRecognitions.add(result);
                            }
                        }

                        tracker.trackResults(mappedRecognitions, currTimestamp);
                        trackingOverlay.postInvalidate();

                        computingDetection = false;


                        /**
                         新建一个操作 ui 的线程
                         * */
                        runOnUiThread(
                                new Runnable() {
                                    @Override
                                    public void run() {
                                        showFrameInfo(previewWidth + "x" + previewHeight);
//                                        showCropInfo(cropCopyBitmap.getWidth() + "x" + cropCopyBitmap.getHeight());
                                        showInference(lastProcessingTimeMs + "ms");

                                        // 更新 RecycleView
                                        updateRecyclerView(results);
                                    }
                                });
                    }
                });
    }

    @Override
    protected int getLayoutId() {
        return R.layout.od_tfe_od_camera_connection_fragment_tracking;
    }


    /**
     * 获取当前展示图像窗口的大小；在 setFragment() 里面的 new LegacyCameraConnectionFragment 时候调用
     */
    @Override
    protected Size getDesiredPreviewFrameSize() {
        return DESIRED_PREVIEW_SIZE;
    }

    @Override
    public void onClick(View v) {

    }

    // Which detection model to use: by default uses Tensorflow Object Detection API frozen
    // checkpoints.
    private enum DetectorMode {
        TF_OD_API;
    }

}
